#!/usr/bin/env python
"""
AMQP data pipeline driver program
"""
import logging
import optparse
import os
import signal
import sys
import time

from dpipelib import Consumer, Master, Workers
from dpipelib import g_log as g_lib_log

g_component = None
g_prog = os.path.basename(sys.argv[0])

g_log = logging.getLogger("dpipe")
_h = logging.StreamHandler()
_h.setFormatter(logging.Formatter("%(levelname)-6s %(name)s.%(funcName)-20s | %(message)s"))
_h.setLevel(logging.NOTSET)
g_log.addHandler(_h)

def handle_signals():
    for signame in 'SIGTERM', 'SIGINT', 'SIGUSR2':
        signum = getattr(signal, signame)
        signal.signal(signum, on_kill)

def on_kill(signo, frame):
    if g_component:
        g_component.shutdown()
    time.sleep(2)
    sys.exit(0)

def add_common_options(op):
    """Add options common to all components."""
    op.add_option('--host', action='store', dest='host',
                  default="localhost",
                  metavar='HOST', help="host to conenct to (default=localhost)")
    op.add_option('--user', action='store', dest='user',
                  default=None,
                  metavar='NAME', help="user name (default=guest user)")
    op.add_option('--password', action='store', dest='pwd',
                default=None,
                metavar='PASS', help="password (default=none)")
    op.add_option('-v', dest='vb', action='count', default=0,
                  help="More verbosity to stderr")
    
def conn_kw(op):
    """Make a dictionary of AMQP connection keywords."""
    d = { 'host' : op.host }
    if op.user is not None:
        d['userid'] = op.user
    if op.pwd is not None:
        d['password'] = op.pwd
    return d

def create_consumer(a):
    op = optparse.OptionParser(usage="%prog (c)onsumer [options]")
    add_common_options(op)
    op.add_option('--stages', dest='ns', default=-1, type='int',
                  help='num. pipeline stages (required)')
    op.add_option('--tasks', dest='nt', default=-1, type='int',
                  help='num. tasks per worker-process (required)')
    op.add_option('--workers', dest='nw', default=-1, type='int',
                  help='num. workers (required)')
    op.add_option('--procs', dest='np', default=-1, type='int',
                  help='num. processes per worker (required)')
    op.add_option('--file', dest='ofile', default=None, help="output file (default=stdout)")
    options, args = op.parse_args(a)
    if options.ns < 1:
        op.error("--stages is required, must be > 0")
    if options.nt < 1:
        op.error("--tasks is required, must be > 0")
    if options.np < 1:
        op.error("--procs per worker is required, must be > 0")
    if options.nw < 1:
        op.error("--workers is required, must be > 0")
    if options.ofile is None:
        ofile = sys.stdout
    else:
        ofile = options.ofile
    set_verbose(options)
    return Consumer(conn_kw(options), ofile=ofile, num_workers=options.nw, num_tasks=options.nt,
                    num_stages=options.ns, num_procs=options.np)

def create_producer(a):
    op = optparse.OptionParser(usage="%prog (p)roducer [options]")
    add_common_options(op)
    op.add_option('--tasks', dest='nt', default=-1, type='int',
                  help='num. tasks per worker-process (required)')
    op.add_option('--workers', dest='nw', default=-1, type='int',
                  help='num. workers (required)')
    op.add_option('--procs', dest='np', default=-1, type='int',
                  help='num. processes per worker (required)')
    op.add_option('--size', dest='sz', default=1024, type='int',
                  help='message size in bytes (%default)')
    options, args = op.parse_args(a)
    if options.nt < 1:
        op.error("--tasks is required, must be > 0")
    if options.nw < 1:
        op.error("--workers is required, must be > 0")
    if options.np < 1:
        op.error("--procs per worker is required, must be > 0")
    set_verbose(options)
    return Master(conn_kw(options), num_workers=options.nw, num_tasks=options.nt,
                  num_procs=options.np, message_size=options.sz)

def create_workers(a):
    op = optparse.OptionParser(usage="%prog (w)orkers [options]")
    add_common_options(op)
    op.add_option('--stages', dest='ns', default=-1, type='int',
                  help='num. pipeline stages (required)')
    op.add_option('--procs', dest='np', default=-1, type='int',
                  help='num. processes for this worker (required)')
    op.add_option('--delay', dest='slp', default=1.0, type='float',
                  help='seconds to sleep for each item of work (default=%default)')
    op.add_option('--sigma', dest='sig', default=0.0, type='float',
                  help='gauss std dev to randomize sleep times (default=%default)')
    op.add_option('--size', dest='sz', default=1024, type='int',
                  help='message size in bytes (default=%default)')
    op.add_option('--tx', dest='tx', default=False, action='store_true', 
                  help='flag for transactional mode')
    options, args = op.parse_args(a)
    if options.ns < 1:
        op.error("--stages is required, must be > 0")
    if options.np < 1:
        op.error("--procs per instance is required, must be > 0")
    set_verbose(options)
    return Workers(conn_kw(options), num_procs=options.np, num_stages=options.ns,
                   sleep_time=options.slp, sleep_sigma=options.sig, 
                   tx=options.tx, message_size=options.sz)

def set_verbose(options):
    if options.vb == 0:
        g_log.setLevel(logging.WARN)
        g_lib_log.setLevel(logging.WARN)
    elif options.vb == 1:
        g_log.setLevel(logging.INFO)
        g_lib_log.setLevel(logging.INFO)
    else:
        g_log.setLevel(logging.DEBUG)
        g_lib_log.setLevel(logging.DEBUG)
    
def main(cmdline=None):
    global g_component, g_log
    if cmdline is None:
        cmdline = sys.argv[1:]
    if len(cmdline) == 0 or cmdline[0] in ('-h', '--help'):
        sys.stderr.write("Help is only available for a specific component.\n"
                         "consumer     {prog} c --help\n"
                         "producer     {prog} p --help\n"
                         "worker(s)    {prog} w --help\n".format(prog=g_prog))
        return 1
    component = cmdline[0].lower()
    if component in ('c', 'consumer'):
        g_component = create_consumer(cmdline[1:])
    elif component in ('p', 'producer'):
        g_component = create_producer(cmdline[1:])
    elif component in ('w', 'worker', 'workers'):
        g_component = create_workers(cmdline[1:])
    else:
        sys.stderr.write("Unknown component: {0}\n"
                         "Use 'c' for consumer, 'p' for producer, or 'w' for workers\n"
                         .format(component))
        return 2
    handle_signals()
    g_log.info("begin component {0}".format(component))
    g_component.run()
    g_log.info("end component {0}".format(component))
    return 0

if __name__ == '__main__':
    sys.exit(main())
